import pandas as pd
from wordcloud import WordCloud
import matplotlib.pyplot as plt

# 读取 Excel 文件
file_path = 'data_with_county.xlsx'
df = pd.read_excel(file_path)

# 筛选出 mark 列为 1 的行
filtered_df = df[df['mark'] == 1]

# 生成 sub_title 列的词云
sub_title_text = ' '.join(filtered_df['sub_title'].dropna().astype(str))

# 创建词云对象，指定支持中文的字体
font_path = 'simhei.ttf'  # 请确保字体文件存在，或者指定一个系统中存在的中文字体路径
wordcloud_sub_title = WordCloud(width=800, height=400, background_color='white', font_path=font_path).generate(sub_title_text)

# 显示 sub_title 列的词云
plt.figure(figsize=(10, 5))
plt.imshow(wordcloud_sub_title, interpolation='bilinear')
plt.axis('off')
plt.title('Sub Title Word Cloud')
plt.savefig('sub_title_wordcloud.png')
plt.show()

# 生成 country 列的词云
# 假设 country 列已经填充完毕，如果没有填充，可以先运行之前的代码填充 country 列
country_text = ' '.join(filtered_df['county'].dropna().astype(str))

# 创建词云对象，指定支持中文的字体
wordcloud_country = WordCloud(width=800, height=400, background_color='white', font_path=font_path).generate(country_text)

# 显示 country 列的词云
plt.figure(figsize=(10, 5))
plt.imshow(wordcloud_country, interpolation='bilinear')
plt.axis('off')
plt.title('Country Word Cloud')
plt.savefig('country_wordcloud.png')
plt.show()